#
# Copyright (C) 2023, Inria
# GRAPHDECO research group, https://team.inria.fr/graphdeco
# All rights reserved.
#
# This software is free for non-commercial, research and evaluation use 
# under the terms of the LICENSE.md file.
#
# For inquiries contact  george.drettakis@inria.fr
#

import torch
import math
# from diff_gaussian_rasterization import GaussianRasterizationSettings, GaussianRasterizer
from diff_gauss import GaussianRasterizationSettings, GaussianRasterizer
from scene.gaussian_model import GaussianModel
from utils.sh_utils import eval_sh
from utils.graphics_utils import get_plucker_coords
import torch.nn.functional as F
import numpy as np
from scene.cameras import Camera
from gsplat.project_gaussians_2d import project_gaussians_2d
from gsplat.rasterize_sum import rasterize_gaussians_sum
import torchvision
from tqdm import tqdm
from einops import repeat
import os
os.environ['CUDA_LAUNCH_BLOCKING'] = '1'


def get_coord(height, width, device):
    y = torch.linspace(-1, 1, height, device=device)
    x = torch.linspace(-1, 1, width, device=device)
    yy, xx = torch.meshgrid(y, x, indexing="ij")
    
    coords = torch.stack([xx, yy], dim=-1).reshape(-1, 2)
    return coords


def render(viewpoint_camera: Camera, 
           pc : GaussianModel, 
           pipe, 
           bg_color : torch.Tensor, 
           scaling_modifier = 1.0, 
           override_color = None):
    """
    Render the scene. 
    
    Background tensor (bg_color) must be on GPU!
    """
 
    # Create zero tensor. We will use it to make pytorch return gradients of the 2D (screen-space) means
    screenspace_points = torch.zeros_like(pc.get_xyz, dtype=pc.get_xyz.dtype, requires_grad=True, device="cuda") + 0
    try:
        screenspace_points.retain_grad()
    except:
        pass

    # Set up rasterization configuration
    tanfovx = math.tan(viewpoint_camera.FoVx * 0.5)
    tanfovy = math.tan(viewpoint_camera.FoVy * 0.5)
    
    lH, lW = int(viewpoint_camera.h), int(viewpoint_camera.w)
    hH, hW = int(viewpoint_camera.image_height), int(viewpoint_camera.image_width)

    raster_settings = GaussianRasterizationSettings(
        image_height=int(viewpoint_camera.h),
        image_width=int(viewpoint_camera.w),
        # image_height=int(viewpoint_camera.image_height),
        # image_width=int(viewpoint_camera.image_width),
        tanfovx=tanfovx,
        tanfovy=tanfovy,
        bg=bg_color,
        scale_modifier=scaling_modifier,
        viewmatrix=viewpoint_camera.world_view_transform,
        projmatrix=viewpoint_camera.full_proj_transform,
        sh_degree=pc.active_sh_degree,
        campos=viewpoint_camera.camera_center,
        prefiltered=False,
        debug=pipe.debug
    )

    rasterizer = GaussianRasterizer(raster_settings=raster_settings)

    means3D = pc.get_xyz
    means2D = screenspace_points
    opacity = pc.get_opacity

    scales = None
    rotations = None
    cov3D_precomp = None
    if pipe.compute_cov3D_python:
        cov3D_precomp = pc.get_covariance(scaling_modifier)
    else:
        scales = pc.get_scaling
        rotations = pc.get_rotation

    shs = None
    colors_precomp = None
    if override_color is None:
        if pipe.convert_SHs_python:
            shs_view = pc.get_shs.transpose(1, 2).view(-1, 3, (pc.max_sh_degree+1)**2)
            dir_pp = (pc.get_xyz - viewpoint_camera.camera_center.repeat(pc.get_shs.shape[0], 1))
            dir_pp_normalized = dir_pp/dir_pp.norm(dim=1, keepdim=True)
            sh2rgb = eval_sh(pc.active_sh_degree, shs_view, dir_pp_normalized)
            colors_precomp = torch.clamp_min(sh2rgb + 0.5, 0.0) # (N, 3)
        else:
            shs = pc.get_shs
    else:
        colors_precomp = override_color
        
    # alpha_attr = pc.get_alpha
    # offsets_attr = pc.get_offset
    # chol_attr = pc.get_cholesky
    
    # TODO: test whether add offset
    # add_offset = False
    # if add_offset:
    #     feats = torch.cat([alpha_attr, offsets_attr, chol_attr], -1)
    # else:
    #     feats = torch.cat([alpha_attr, chol_attr], -1)
    # rendered_image, rendered_depth, pseudo_normal, rendered_alpha, radii, feat_map = rasterizer(means3D=means3D,
    #                             means2D=means2D,
    #                             shs=None,
    #                             colors_precomp=torch.zeros_like(feats[:, :3]),
    #                             opacities=opacity,
    #                             scales=scales,
    #                             rotations=rotations,
    #                             cov3Ds_precomp=cov3D_precomp,
    #                             extra_attrs=feats)  
    
    # BLOCK_W, BLOCK_H = 16, 16
    # tile_bounds = (int((hW+BLOCK_W-1)//BLOCK_W), int((hH+BLOCK_H-1)//BLOCK_H), 1)
    # background = torch.ones(3).cuda()
    
    # if add_offset:    
    #     opacity, offset, chol_weights = torch.split(feat_map, (1, 2, 3), dim=0) # (3, H, W)
    #     offset = offset.permute(1,2,0).reshape(-1, 2)                       # (H*W, 2)
    #     coords = get_coord(H, W, rendered_image.device).detach()     # (-1, 1)
    #     xyz1 = coords[:,0:1] + (pc.config.low_ratio / 2) * (1 + offset[:,0:1]) / W
    #     xyz2 = coords[:,1:2] + (pc.config.low_ratio / 2) * (1 + offset[:,1:2]) / H
    #     coords = torch.cat((xyz1, xyz2), dim = 1)
    # else:
    #     opacity, chol_weights = torch.split(feat_map, (1, 3), dim=0) # (3, H, W)
    #     coords = get_coord(H, W, rendered_image.device).detach()     # (-1, 1)
    #     # xyz1 = coords[:,0:1] + pc.config.low_ratio / 2 / W
    #     # xyz2 = coords[:,1:2] + pc.config.low_ratio / 2 / H
    #     # coords = torch.cat((xyz1, xyz2), dim = 1)
    
    # color = rendered_image.permute(1,2,0).reshape(-1, 3)
    # opacity = opacity.permute(1,2,0).reshape(-1, 1)                     # (H*W, 1)
    # chol_weights = chol_weights.permute(1,2,0).reshape(-1, 3)   
    # chol = pc.config.low_ratio / 2 * (chol_weights * pc.chol_min[None] + (1-chol_weights) * pc.chol_max)  # (H*W, 3)
    
    BLOCK_W, BLOCK_H = 16, 16
    tile_bounds = (int((hW+BLOCK_W-1)//BLOCK_W), int((hH+BLOCK_H-1)//BLOCK_H), 1)
    background = torch.ones(3).cuda()
    
    ob_view = means3D - viewpoint_camera.camera_center  # (N, 3)
    ob_dist = ob_view.norm(dim=1, keepdim=True)        # (N, 1)
    ob_view = ob_view / ob_dist  
    
    feats = pc.get_feats
    rendered_view, _, _, _, radii, feat_map = rasterizer(
                                means3D=means3D,
                                means2D=means2D,
                                shs=None,
                                colors_precomp=ob_view,
                                opacities=opacity,
                                scales=scales,
                                rotations=rotations,
                                cov3Ds_precomp=cov3D_precomp,
                                extra_attrs=feats)  
    
    view_prior = rendered_view.permute(1,2,0).reshape(-1, 3)
    view_feats = feat_map.permute(1,2,0).reshape(-1, pc.feat_dim)   # (lH*lW / (ratio**2), feat_dim)
    view_feats = torch.cat([view_prior, view_feats], -1)
    
    if pc.add_plucker:
        plucker_coords = viewpoint_camera.plucker_rays.cuda()       # (lH*lW / (ratio**2), 6)
        view_feats = torch.cat([view_feats, plucker_coords], -1) 
    
    offsets = pc.get_offset_mlp(view_feats).reshape(-1, 2)
    chol_weights = pc.get_cholesky_mlp(view_feats).reshape(-1, 3)
    opacities = pc.get_opacity_mlp(view_feats).reshape(-1)
    colors = pc.get_color_mlp(view_feats).reshape(-1, 3)
    
    # chol = pc.config.low_ratio / 2 * (chol_weights * pc.chol_min[None] + (1-chol_weights) * pc.chol_max)  # (H*W, 3)
    chol = pc.config.low_ratio / 2 * (chol_weights * pc.chol_min[None] + (1-chol_weights) * pc.chol_max[None])  # (H*W, 3)
    
    coords = get_coord(lH, lW, offsets.device).detach()              # (lH*lW, 2)
    # np.save("init_coords", coords.detach().cpu().numpy())
    repeat_coords = repeat(coords, 'n (c) -> (n k) (c)', k=pc.n_offsets)
    xyz1 = repeat_coords[:,0:1] + pc.config.low_ratio * 2 * offsets[:,0:1] / lW
    xyz2 = repeat_coords[:,1:2] + pc.config.low_ratio * 2 * offsets[:,1:2] / lH
    coords = torch.cat((xyz1, xyz2), dim = 1)
    
    mask = opacities > 0.0
    
    coords = coords[mask]
    opacity = opacities[mask][:, None]
    colors = colors[mask]
    chol = chol[mask]
            
    xys, depths, radii_, conics, num_tiles_hit = project_gaussians_2d(coords, chol, hH, hW, tile_bounds)
    
    out_img = rasterize_gaussians_sum(xys, depths, radii_, conics, num_tiles_hit,
                colors, opacity, hH, hW, BLOCK_H, BLOCK_W, background=background, return_alpha=False)
    
    # np.save("coords", coords.detach().cpu().numpy())
    # np.save("xys", xys.detach().cpu().numpy())
    # breakpoint()
    
    rendered_image = out_img.permute(2,0,1)
    
    return {"render": rendered_image,
            "viewspace_points": screenspace_points,
            "visibility_filter" : radii > 0,
            "radii": radii}

